"""
author: liuxu
date: 2025/8/13
description: 
"""
import pandas as pd
import paramiko
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from socket import timeout as SocketTimeout


def connect_and_execute(row, command):
    """连接单个节点并执行命令"""
    ip = row.get("节点ip")
    username = row.get("用户名")
    password = row.get("密码")

    # 检查必要字段是否为空
    if not all([ip, username, password]):
        return f"节点配置不完整: IP={ip}, 用户名={username}"

    result = {"ip": ip, "output": "", "error": "", "status": "success"}

    try:
        # 创建 SSH 客户端
        with paramiko.SSHClient() as client:
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            # 建立 SSH 连接（增加 banner_timeout）
            client.connect(
                hostname=ip,
                username=username,
                password=password,
                timeout=10,
                banner_timeout=20
            )

            # 执行命令（增加 get_pty=True 解决某些环境问题）
            stdin, stdout, stderr = client.exec_command(command, get_pty=True)

            # 等待命令完成（最多等待60秒）
            exit_status = stdout.channel.recv_exit_status()

            # 读取输出
            out = stdout.read().decode("utf-8", errors="ignore").strip()
            err = stderr.read().decode("utf-8", errors="ignore").strip()

            if exit_status != 0:
                result["status"] = f"命令失败 (exit code {exit_status})"
                result["error"] = err
            else:
                result["output"] = out

    except paramiko.AuthenticationException:
        result["status"] = "认证失败"
    except paramiko.SSHException as e:
        result["status"] = f"SSH错误: {str(e)}"
    except SocketTimeout:
        result["status"] = "连接超时"
    except Exception as e:
        result["status"] = f"未知错误: {str(e)}"

    return result


def execute_command_on_nodes(excel_path, command, max_workers=5):
    """批量在节点上执行命令"""
    # 读取 Excel 节点信息
    try:
        df = pd.read_excel(excel_path)
    except Exception as e:
        print(f"读取Excel失败: {str(e)}")
        return

    # 筛选需要测试的节点
    test_nodes = df[df["测试"].astype(str).str.lower() == "yes"]

    if test_nodes.empty:
        print("没有找到需要测试的节点")
        return

    print(f"将在 {len(test_nodes)} 个节点上执行命令: {command}")

    # 使用线程池并行执行
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {}
        for index, row in test_nodes.iterrows():
            futures[executor.submit(connect_and_execute, row, command)] = row["节点ip"]

        # 处理结果
        for future in as_completed(futures):
            ip = futures[future]
            try:
                result = future.result()
                print(f"\n==== {ip} ====")
                print(f"状态: {result['status']}")

                if result["error"]:
                    print(f"错误信息:\n{result['error']}")

                if result["output"]:
                    print(f"命令输出:\n{result['output']}")

                print("-" * 40)
            except Exception as e:
                print(f"处理 {ip} 结果时出错: {str(e)}")


if __name__ == "__main__":
    # 解析命令行参数
    if len(sys.argv) < 2:
        print("用法: python batch_exec.py <Excel文件路径> <要执行的命令>")
        print("可选参数: -w <线程数> (默认5)")
        sys.exit(1)

    excel_path = sys.argv[1]
    command = sys.argv[2]

    # 解析线程数参数
    max_workers = 5
    if "-w" in sys.argv:
        try:
            idx = sys.argv.index("-w")
            max_workers = int(sys.argv[idx + 1])
        except (ValueError, IndexError):
            print("警告: 无效的线程数参数，使用默认值5")

    execute_command_on_nodes(excel_path, command, max_workers)